import json
import logging
import os
import re
from collections import deque

from sonetto.agents.base import BaseAgent, ContextLevel
from sonetto.agents.chat import ChatAgent
from sonetto.agents.tool import ToolAgent
from sonetto.core.message import Message
from sonetto.core.routing.broadcast import CommonBroadcastConditions
from sonetto.core.scheduler import Scheduler
from sonetto.core.toolset import McpToolset, field

base_url_glm = "https://open.bigmodel.cn/api/paas/v4"
api_key_glm = "ff37fbba1e80d4cb5f38ac819893f4a1.c6T0Zm3Sd9YODRLV"
model_glm = "glm-4-flash"

base_url_ds = "https://api.deepseek.com"
api_key_ds = "sk-3bb4d61950df4f4d833f800848e1f461"
model_ds = "deepseek-chat"

toolset = McpToolset()


class CoordinatorAgent(BaseAgent):

    def __init__(self, agent_id: str, working_directory: str):
        super().__init__(agent_id)
        self.directory = working_directory
        self.document_queue = deque([
            doc for doc in os.listdir(working_directory) if doc.endswith('md')
        ])
        logging.debug(f"文档队列：{self.document_queue}")
        self.current_document = None
        self.current_output = None
        self.current_line = 0
        self.req_line, self.req_end = 0, 0
        self.total_lines = 0
        self.window_size = 100
        self.request_received, self.check = False, False
        self.json_obj = {}

    def list_directory(self):
        return os.listdir(self.directory)

    def pretty_print(self, content: dict[str], lineno: bool = True) -> str:
        ret = ""
        for key, value in content.items():
            if lineno:
                ret += f"第{key}行: {value}\n"
            else:
                ret += f"{value}\n"
        ret += "-----------------\n"
        return ret

    def read_lines(self, start: int, end: int, lineno: bool = True) -> list[str]:
        """
        读取文件中指定范围的行
        
        Args:
            filename: 文件路径
            start: 起始行号（从1开始，包含）
            end: 结束行号（从1开始，包含）
            
        Returns:
            包含行号和内容的字符串列表，每个字符串格式为 "行号: 内容"
            
        Raises:
            FileNotFoundError: 文件不存在时抛出
            ValueError: 起始或结束行号无效时抛出
            TypeError: 行号不是整数时抛出
        """
        # 参数验证
        if not isinstance(start, int) or not isinstance(end, int):
            raise TypeError("行号必须是整数")

        if start < 1:
            raise ValueError("起始行号必须大于等于1")

        if end < start:
            raise ValueError("结束行号不能小于起始行号")

        try:
            result = {}
            with open(self.current_document, 'r', encoding='utf-8') as file:
                line_num = 0
                for line in file:
                    line_num += 1

                    if line_num < start:
                        continue

                    if line_num <= end:
                        result[line_num] = f"{line.rstrip()}"
                    else:
                        break

                if not result and start > line_num:
                    raise ValueError(f"文件只有{line_num}行，少于起始行号{start}")
                return self.pretty_print(result, lineno)

        except FileNotFoundError:
            raise FileNotFoundError(f"文件 '{self.current_document}' 不存在")

    def jsonl_push_with_lineno(self, index: str, start: int, end: int) -> None:
        """
        读取 markdown 文件中指定范围的行，然后将索引和内容写入到 JSONL 文件中
        
        Args:
            index: 文本块摘要
            start: 起始行号（从1开始，包含）
            end: 结束行号（从1开始，包含）
            
        Raises:
            FileNotFoundError: 文件不存在时抛出
            ValueError: 起始或结束行号无效时抛出
            TypeError: 行号不是整数时抛出
            IOError: 文件无法访问或写入时抛出
        """
        start = int(start)
        end = int(end)
        # 参数验证
        if not isinstance(start, int) or not isinstance(end, int):
            raise TypeError("行号必须是整数")

        if start < 1:
            raise ValueError("起始行号必须大于等于1")

        if end < start:
            raise ValueError("结束行号不能小于起始行号")

        try:
            result = ''
            with open(self.current_document, 'r', encoding='utf-8') as file:
                line_num = 0
                for line in file:
                    line_num += 1

                    if line_num < start:
                        continue

                    if line_num <= end:
                        result += f"{line.rstrip()}\n"
                    else:
                        break

                if not result and start > line_num:
                    raise ValueError(f"文件只有{line_num}行，少于起始行号{start}")

            # 构造要追加的 JSON 对象
            entry = {"index": index, "content": result}
            self.json_obj = entry
            self.req_end = end + 1
            self.check = True

            return f"""你的分块请求已被用户收到。请你等候用户的审查结果。"""

        except FileNotFoundError as e:
            logging.error(e)
            # exit()
            raise FileNotFoundError(f"文件 '{self.current_document}' 不存在")
        except Exception as e:
            raise IOError(f"向 {self.current_output} 追加内容时出错: {str(e)}")

    def request_more(self) -> None:
        self.request_received = True
        return "这是一条通知，告诉你请求已发出。你无需回复，等候用户提问即可。"

    def request_finished(self) -> None:
        self.request_received = False
        self.req_line += self.window_size
        logging.debug(self.req_line)

    def count_lines(self) -> int:
        try:
            with open(self.current_document, 'r', encoding='utf-8') as file:
                return len(file.readlines())
        except FileNotFoundError:
            raise FileNotFoundError(f"文件 '{self.current_document}' 不存在")

    def turn2next_document(self) -> None:
        if self.document_queue:
            self.current_document = os.path.join(self.directory,
                                                 self.document_queue.popleft())
            self.current_output = f"{self.current_document}.jsonl"
            self.current_line, self.req_line = 1, 1
            self.total_lines = self.count_lines()
            return True
        else:
            return False

    def approve(self) -> str:
        # 检查文件是否存在，不存在则创建
        file_exists = os.path.isfile(self.current_output)
        if not file_exists:
            with open(self.current_output, 'w', encoding='utf-8'):
                pass
        file_size = os.path.getsize(self.current_output) if file_exists else 0

        with open(self.current_output, 'a', encoding='utf-8') as file:
            # 如果文件已存在且非空，确保新条目从新行开始
            if file_exists and file_size > 0:
                # 检查文件是否以换行符结束
                with open(self.current_output, 'rb') as check_file:
                    if file_size > 0:
                        check_file.seek(file_size - 1)
                        last_char = check_file.read(1)
                        if last_char != b'\n':
                            file.write('\n')

            # 将 JSON 对象转换为字符串并写入文件
            json_str = json.dumps(self.json_obj, ensure_ascii=False)
            file.write(json_str + '\n')
            self.current_line = self.req_end + 1
            self.check = False
            self.json_obj = {}

            return "你的批准已经被后台记录，收到请回复。"

    def forward(self, message: Message) -> Message:
        # Logic for orchestrating the entire workflow
        # Tracking current document and processing stage
        # Routing messages between agents

        sender = message.sender
        receiver = message.receiver
        content = message.content

        if sender == '__init__':
            if not self.turn2next_document():
                return Message(content="没有文档需要处理",
                               sender=self.agent_id,
                               receiver=sender)
            else:
                content = f"""
接下来开始对文档 '{self.current_document}' 进行分块。其从第{self.req_line}行开始的{self.window_size}行内容如下所示：
{self.read_lines(self.req_line, self.req_line + self.window_size - 1)}
请确认内容中是否包含完整的文本块，如果是，请调用 `jsonl_push_with_lineno` 接口，否则请调用 `request_more` 接口请求更多文档内容。
"""
                self.req_line += self.window_size
                return Message(content=content,
                               sender=self.agent_id,
                               receiver='chunker')

        # 检查是否有文档需要处理
        if self.current_line >= self.total_lines:
            if not self.turn2next_document():
                return Message(content="所有文档均已整理完毕，圆满结束！！！",
                               sender=self.agent_id,
                               receiver=sender)
            else:
                content = f"""
接下来开始对文档 '{self.current_document}' 进行分块。其从第{self.req_line}行开始的{self.window_size}行内容如下所示：
{self.read_lines(self.req_line, self.req_line + self.window_size - 1)}
请确认尚未分块的内容中是否包含完整的文本块，如果是，请调用 `jsonl_push_with_lineno` 接口，否则请调用 `request_more` 接口请求更多文档内容。
"""
                self.request_finished()
                return Message(content=content,
                               sender=self.agent_id,
                               receiver=sender)

        if self.request_received:  # 如果智能体请求查看更多内容
            content = f"""
当前分块进度：
- 当前文档：{self.current_document}
- 已完成分块的行数：{self.current_line - 1}
- 总行数：{self.total_lines}

我收到了你对更多文档内容的请求。
文档 '{self.current_document}' 从第{self.req_line}行开始的{self.window_size}行内容如下所示：
```
{self.read_lines(self.req_line, self.req_line + self.window_size - 1)}
```
请确认第{self.current_line}行和第{self.req_line + 99}行之间是否包含语义完整的文本块，
如果是，请调用 `jsonl_push_with_lineno` 接口，
否则请调用 `request_more` 接口请求更多文档内容。

"""
            self.request_finished()
            return Message(content=content,
                           sender=self.agent_id,
                           receiver=sender)
        elif self.check and sender == 'chunker':  # 如果智能体请求审查
            return Message(content=f"""
请帮我检查以下文本块是否符合分块要求：
- 文本块摘要：{self.json_obj['index']}
- 文本块内容：
```markdown
{self.json_obj['content']}
```

该文本块在文档中后面接着的部分内容是：
```markdown
{self.read_lines(self.req_end, self.req_end + 21, False)}
```

如果符合要求，您只需调用 `approve_chunk` 接口，除此之外无需多言。
只有在文本块不符要求时才需要您给出边界放置的具体反馈。
""",    # 如果符合要求，请调用 `approve_chunk` 接口，否则请提出修改意见。
                           sender=self.agent_id,
                           receiver='checker')
        elif sender == 'checker' and self.check:  # 如果审查员审查不通过
            self.check = False
            return Message(content=content, sender=sender, receiver='chunker')
        else:
            content = f"""
你的上一个分块审核已经通过。

当前分块进度：
- 当前文档：{self.current_document}
- 已完成分块的行数：{self.current_line - 1}
- 总行数：{self.total_lines}

请确认第{self.current_line}行和第{self.req_line - 1}行之间是否包含语义完整的文本块，
如果是，请调用 `jsonl_push_with_lineno` 接口。
否则请调用 `request_more` 接口请求更多文档内容。

"""
            return Message(content=content,
                           sender=self.agent_id,
                           receiver='chunker')


coordinator = CoordinatorAgent(agent_id="coordinator",
                               working_directory=".\docs")


@toolset.register_tool(name="Jsonl_Push_With_Lineno",
                       description="将文本块摘要和指定范围的行送交审核")
def jsonl_push_with_lineno(index: str = field(description="文本块摘要"),
                           start: int = field(description="文本块在文档中的起始行号"),
                           end: int = field(
                               description="文本块在文档中的末尾行号（包括该行）")) -> None:
    return coordinator.jsonl_push_with_lineno(index, start, end)


@toolset.register_tool(name="Request_More", description="请求更多文档内容")
def request_more() -> str:
    return coordinator.request_more()


chunker = ToolAgent(agent_id="chunker",
                    base_url=base_url_ds,
                    api_key=api_key_ds,
                    model=model_ds,
                    toolset=toolset,
                    system_prompt=''''
你是一个文档分块专家，你的首要任务是识别文档中文本的语义边界，
生成文本块摘要，并调用接口将分块结果写入文件，用于后续的知识工程搭建。
用户会给你反馈分块进度和审查结果，你需要根据用户的反馈及时调整分块策略。

### 语义分块指南：
1. 每个块应在语义上连贯且独立。
2. 直到出现了明确的语义边界后再进行分块：
    - 分节边界（章节、标题）
    - 主题转换
    - 形成逻辑单元的段落组
3. 对于图片或表格，就近分块，不要将图片或表格与其描述分开。
4. 分块摘要应当与文本块的主题相关，不要出现“封面信息”、“第一章”、“第二节”这种无意义的摘要。
5. 不要对已通过审核的文本块重新分块，除非用户要求修改。

### 对于每一个文本块，提供以下信息：
1. 文本块摘要：分级简短描述文本块的主题，例如：龙腾集团员工手册-员工权益-年假
2. 文本块范围：文本块在文档中的起始行号和结束行号（行号从 1 开始）

由于文档非常庞大，用户只能分多次为你提供文档内容，每次只能提供100行。
在你收到的没有分块的内容中，如果你识别到了语义完整的文本块，请调用 `jsonl_push_with_lineno` 接口，
将文本块摘要和指定范围的行送交审核。
否则，你应当先记下收到的文档内容，并调用 `request_more` 接口请求更多文档内容。
用户在收到请求后，会在下一次提问中提供更多文档内容。

一次回答中只能调用一次接口。

''')

toolset_checker = McpToolset()


@toolset_checker.register_tool(name="Approve_Chunk", description="批准文本块，并由后台记录下此次批准")
def approve() -> str:
    return coordinator.approve()


checker = ToolAgent(agent_id="checker",
                    base_url=base_url_ds,
                    api_key=api_key_ds,
                    model=model_ds,
                    toolset=toolset_checker,
                    system_prompt='''
你是一位严格的语义分析专家，善于评估文本块的质量。和你对话的是一个文档分块员工。

你的工作是验证提交的文本块是否保持语义完整性并代表完整的知识单元。

### 评估标准：
1. 语义完整性：文本块是否包含理解其内容所需的所有必要上下文？
2. 边界质量：文本块边界是否位于自然的语义过渡处？结合下文，判断边界是否需要后置。
比如员工有可能会在目录中间划分边界，这时候边界就需要后置。
3. 信息完整性：文本块语义上是否存在严重的割裂？
4. 自包含性：文本块是否可以独立理解？
5. 信息重复：文本块是否重复了之前文本块的信息？
6. 图片和表格就近分块

你应当尽量严格地评估文本块的质量，确保文本块的质量符合要求。

如果文本块通过审核，请直接调用工具API：`approve_chunk`，不要附加任何额外反馈。
如果不通过，请用中文提供反馈并要求修改。

对于不合格的文本块，你需要提供关于边界放置的具体反馈：文本块边界应当前置还是后置？
如果需要前置，告诉用户前置的原因和位置；如果需要后置，提示员工往后翻文档内容即可（员工看到的上下文比你更少，所以他们的失误很有可能是这个原因）。

''')

# `request_more` 的调用者只能是你！请不要让用户调用这个接口，也不要复述用户的请求。
scheduler = Scheduler()
scheduler.register_agent(coordinator)
scheduler.register_agent(chunker)
scheduler.register_agent(checker)
# scheduler.set_route(coordinator.agent_id, chunker.agent_id)
scheduler.set_route(chunker.agent_id, coordinator.agent_id)
scheduler.set_route(checker.agent_id, coordinator.agent_id)

if __name__ == "__main__":
    logging.debug("文档分块任务开始")
    scheduler.run_stream("请完成文档分块任务",
                         cancel_token="圆满结束",
                         cancel_from=[coordinator.agent_id],
                         max_step=114)
    logging.debug("文档分块任务结束")
    # print([json.loads(ele['content']) for ele in chunker.get_context(ContextLevel.SESSION)])
'''
如仍有完整分块需要写入，请继续调用 `jsonl_push_with_lineno` 接口。
如果需要查看更多文档内容，请调用 `request_more` 接口。
'''
